
from fastapi import FastAPI
from datasource.obtains import engine, get_async_session, sync_pool, redis_pool, get_async_redis
from model import get_sync_db, Status, to_json, create_table_all, sync_engine
from datetime import datetime, timedelta
from typing import Any, Union, Dict
from jose import jwt
{{if eq .isUser true}}from starlette.responses import JSONResponse{{end}}
{{if eq .isUser true}}from variable import ACCESS_TOKEN_EXPIRE_MINUTES, SECRET_KEY, ALGORITHM{{end}}
from apscheduler.schedulers.background import BackgroundScheduler
from pytz import utc
import variable
from apscheduler.executors.pool import ProcessPoolExecutor
from multiprocessing import cpu_count
import os

{{if eq .isUser true}}
# 解析token装饰器
def analysis_token_decorator(func):
    async def wrapper(session, token):
        token_str_list = await verify_token_format(token)
        payload = await analysis_token(token_str_list[1])
        return await func(session, payload["sub"])

    return wrapper
{{- end}}

# 事务处理装饰器
def transaction_decorator(func):
    # async def wrapper(self, session, *params, **kwparams):
    async def wrapper(session, *params):
        try:
            # await func(self, session, *params, **kwparams)
            await func(session, *params)

        except Exception as e:
            await session.rollback()
            await session.close()

            return {"message": "数据库操作失败", "code": 500, "errors": repr(e)}
        else:
            await session.commit()

        finally:
            await session.close()

        return {"message": "操作成功", "code": 200, "data": None}

    return wrapper

# 事务处理装饰器同步
def transaction_decorator_sync(func):
    def wrapper(session, *params):
        try:
            func(session, *params)

        except Exception as e:
            session.rollback()
            session.close()

            return {"message": "数据库操作失败", "code": 500, "errors": repr(e)}
        else:
            session.commit()

        finally:
            session.close()

        return {"message": "操作成功", "code": 200, "data": None}

    return wrapper

def start_task():
    cpu = cpu_count()

    executors = {
        'default': {
            'type': 'threadpool',
            'max_workers': cpu/2
        },
        'processpool': ProcessPoolExecutor(max_workers=cpu/4)
    }
    variable.TASK = BackgroundScheduler(executors=executors, timezone="Asia/Shanghai")
    variable.TASK.start()

def register_scheduler(app: FastAPI):
    @app.on_event("startup")
    def start_event():
        create_table_all()
        start_task()

    @app.on_event("startup")
    async def start_event():
        # app.state.db = await get_async_session()
        # app.state.redis = await get_async_redis()

        pass

    @app.on_event("shutdown")
    def shutdown_event():
        sync_engine.dispose()
        sync_pool.disconnect()

        variable.TASK.shutdown()

    @app.on_event("shutdown")
    async def shutdown_event():
        await engine.dispose()
        await redis_pool.disconnect()

{{- if eq .isUser true}}
# 生成token
async def create_access_token(subject: Union[str, Any], expires_delta: timedelta = None) -> Dict[str, Any]:
    """
    # 生成token
    :param subject: 保存到token的值
    :param expires_delta: 过期时间
    :return:
    """
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + datetime.timedelta(
            minutes=ACCESS_TOKEN_EXPIRE_MINUTES
        )
    to_encode = {"exp": expire, "sub": str(subject)}
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return {"token": encoded_jwt, "expires_delta": expires_delta}

# 解析token
async def analysis_token(token):
    try:
        payload = jwt.decode(
            token,
            SECRET_KEY, algorithms=ALGORITHM
        )
        return payload
    except (jwt.JWTError, jwt.ExpiredSignatureError, AttributeError):
        # 抛出自定义异常， 然后捕获统一响应
        return False

async def verify_token_format(tokenStr):
    if tokenStr is None or tokenStr == "":
        return JSONResponse(content={"code": 604, "data": None, "message": "token不能为空"})

    tokenList = tokenStr.split(" ", 2)
    if len(tokenList) != 2:
        return JSONResponse(content={"code": 604, "data": None, "message": "token格式错误"})

    if tokenList[0] != "Bearer":
        return JSONResponse(content={"code": 604, "data": None, "message": "token头错误"})

    return tokenList
{{- end}}

def command_line(cmd: str):
    contents = os.popen(cmd)
    print(contents.read())

